from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.responses import FileResponse
import shutil
import os
import tempfile
from rag_chat import load_vectorstore, process_document_with_requirements

app = FastAPI()

@app.post("/process_document/")
async def process_document(file: UploadFile = File(...), question: str = Form(...)):
    """
    处理上传的文档，根据问题使用RAG技术修改文档格式并返回修改后的文档
    
    Args:
        file: 上传的原始文档文件
        question: 用户问题
        
    Returns:
        修改后的文档文件
    """
    # 创建临时目录用于存储上传的文件
    with tempfile.TemporaryDirectory() as temp_dir:
        # 保存上传的文件到临时路径
        temp_input_path = os.path.join(temp_dir, file.filename)
        with open(temp_input_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        
        try:
            # 加载向量数据库
            vectorstore = load_vectorstore("Rag/faiss_index")
            
            # 处理文档，直接使用上传的文件
            success = process_document_with_requirements(question, vectorstore, temp_input_path)
            
            if success:
                # 返回修改后的文档
                output_doc_path = os.path.join("Output", "modified_example.docx")
                if os.path.exists(output_doc_path):
                    return FileResponse(
                        output_doc_path, 
                        media_type='application/vnd.openxmlformats-officedocument.wordprocessingml.document',
                        filename="modified_document.docx"
                    )
                else:
                    raise HTTPException(status_code=500, detail="处理后的文件未找到")
            else:
                raise HTTPException(status_code=500, detail="文档处理失败")
                
        except FileNotFoundError as e:
            raise HTTPException(status_code=404, detail=f"向量数据库未找到: {str(e)}")
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"处理过程中发生错误: {str(e)}")

@app.get("/")
def read_root():
    return {"message": "文档处理 API 服务正在运行"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)